package com.agentframework.core.plan;

import com.agentframework.core.agent.Agent;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.prompt.AssistantPromptTemplate;
import org.springframework.ai.chat.prompt.PromptTemplate;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * @author wangjun
 **/
public class Planner {

    protected static final Logger logger = LoggerFactory.getLogger(Planner.class);

    String STRUCTURAL_CONTEXT = """
    ## 用户需求
    {user_requirement}
    
    ## 上下文
    {context}
    
    ## 当前计划
    {tasks}
    
    ## 当前任务
    {current_task}
    """;

    String PLAN_STATUS = """
    ## 任务目标
    {goal}
    ## 已完成的任务
    {finish_tasks}
    
    ### 任务执行结果
    {task_results}
 
    ## 当前任务
    {current_task}
    
    ## 任务执行规则
    你需要直接执行当前任务 {current_task}，并调用当前任务提供的工具 {tools}.
    已经完成了的任务不需要再告诉我已经完成，必须在当前消息中直接调用该工具
    
    """;

    Plan plan;

    int maxRetries = 3;

    List<Message> workingMemory = new ArrayList<>();

    boolean autoRun = true;

    public Planner() {
        this.plan = new Plan();
    }

    public void createPlan(Agent agent, String message) {

        try{
            agent.getContext().setTodo(((PlanAgent) agent).writePlanAction);
            this.plan.setGoal(message);

            boolean planConfirmed = false;

            while (!planConfirmed) {
                List<Message> usefulMemories = getUsefulMemories();
                String resp = agent.getContext().getTodo().execute(usefulMemories).getText();
                this.workingMemory.add(new AssistantPromptTemplate(resp).createMessage());
                GenerPlanResult result = generatePlanFromResp(resp);

                if (!result.isPlanValid && maxRetries > 0) {
                    String errorMsg = String.format("生成的计划无效，错误为：%s，请尝试重新生成，记住只生成整个计划或单个更改的任务", result.error);
                    this.workingMemory.add(new AssistantPromptTemplate(errorMsg).createMessage());
                    maxRetries --;
                    continue;
                }

                AskReviewResult reviewResult = askReview(ReviewConst.TASK_REVIEW_TRIGGER);
                planConfirmed = reviewResult.planConfirmed;

            }

            this.workingMemory.clear();

        } finally {
            agent.getContext().setTodo(null);
        }
    }

    public void processTaskResult(Agent agent, TaskResult taskResult) {
        AskReviewResult result = askReview(taskResult, ReviewConst.TASK_REVIEW_TRIGGER);
        if (result.planConfirmed) {
            confirmTask(agent, getPlan().getCurrentTask(), taskResult, result.review);
        } else if (result.review.contains("redo")) {
            // Ask the Agent to redo this task with help of review feedback,
            // useful when the tool run is successful but the procedure or result is not what we want
            // simply pass, not confirming the result
        } else {
            createPlan(agent, getUsefulMemories().get(0).getText());
        }
    }

    public AskReviewResult askReview(TaskResult taskResult, String trigger) {
        if (!autoRun) {
            return new AskReview(this.workingMemory, this.plan, trigger).execute();
        }

        return new AskReviewResult("", taskResult != null ? taskResult.isSuccess() : true);
    }

    public AskReviewResult askReview(String trigger) {
        return askReview(null, trigger);
    }

    public record AskReviewResult(String review, boolean planConfirmed) {

    }

    public void confirmTask(Agent agent, Task task, TaskResult taskResult, String review) {
        task.updateTaskResult(taskResult);
        getPlan().finishCurrentTask();
        workingMemory.clear();

        boolean confirmed = ReviewConst.CONTINUE_WORDS.contains(review.toLowerCase()) ||
                review.toLowerCase().contains(ReviewConst.CONTINUE_WORDS.get(0));

        if (confirmed) {
            Message message = new PromptTemplate(review).createMessage();
            workingMemory.add(message);
            createPlan(agent, getUsefulMemories().get(0).getText());
        }

    }

    public List<Message> getUsefulMemories() {
        String userRequirement = plan.getGoal();
        String context = plan.getContext() == null ? "" : plan.getContext();
        String tasks = plan.getTasks() == null ? "[]" : plan.getTasks().stream().map(Task::toString).collect(Collectors.joining( "\n\n"));
        String currentTask = plan.getCurrentTask() == null ? "" : plan.getCurrentTask().toString();
        workingMemory.add(PromptTemplate.builder().template(STRUCTURAL_CONTEXT).variables(Map.of("user_requirement", userRequirement,
                "context", context, "tasks", tasks, "current_task", currentTask)).build().createMessage());
        return workingMemory;
    }

    public String getPlanStatus() {
        List<Task> finishedTasks = getPlan().getFinishedTasks();
        String finishTasks = finishedTasks.stream().map(Task::toString).collect(Collectors.joining( "\n\n"));
        String taskResults = finishedTasks.stream().map(Task::getResult).collect(Collectors.joining( "\n\n"));
        List<String> tools = getPlan().getCurrentTask().getTaskTools();
        String toolsDesc = String.join( "\n\n", tools);

        return PromptTemplate.builder().template(PLAN_STATUS).variables(Map.of(
                "goal", this.getPlan().getGoal(),
                "finish_tasks", finishTasks,
                "task_results", taskResults,
                "current_task", getPlan().getCurrentTask().toString(),
                "tools", toolsDesc)).build().create().getContents();
    }

    public Plan getPlan() {
        return this.plan;
    }

    public void setAutoRun(boolean autoRun) {
        this.autoRun = autoRun;
    }

    private GenerPlanResult generatePlanFromResp(String resp) {

        String regex = "```json([\\s\\S]*?)```";

        Pattern pattern = Pattern.compile(regex);
        Matcher matcher = pattern.matcher(resp);

        if (matcher.find()) {
            String content = matcher.group(1);
            resp = content.replace("(", "{").replace(")", "}");
        }

        try {
            List<Task> tasks = new ObjectMapper().readValue(resp, new TypeReference<>() {});

            if (tasks.size() == 1 || !tasks.get(0).getDependentTaskIds().isEmpty()) {
                if (!tasks.get(0).getDependentTaskIds().isEmpty() && tasks.size() > 1) {
                    // tasks[0].dependent_task_ids means the generated tasks are not a complete plan
                    // for they depend on tasks in the current plan, in this case, we only support updating one task each time
                    logger.warn("Current plan will take only the first generated task if the generated tasks are not a complete plan");
                }

                Task task = tasks.get(0);

                if (plan.hasTaskId(task.getTaskId())) {
                    plan.replaceTask(task);
                } else {
                    plan.appendTask(task);
                }

            } else {
                plan.addTasks(tasks);
            }

        } catch (JsonProcessingException e) {
            return new GenerPlanResult(false, e.getMessage());
        }
        return new GenerPlanResult(true);
    }

    record GenerPlanResult(boolean isPlanValid, String error) {
        public GenerPlanResult(boolean isPlanValid) {
            this(isPlanValid, "");
        }
    }

}
